[OMNIML-6071] Calibrate a distributed model on per-rank data - #2740

What does this PR do?

Type of change: New feature

The pieces a calibration loop needs once the model is spread across ranks.

• shard_dataloader gives each rank its own shard of the data.
  • Map-style datasets use a DistributedSampler with drop_last off, so every rank runs the same number of steps; a rank that exits early hangs the others on the per-step collectives.
  • An IterableDataset gets a strided copy of the stream, which must yield the same samples in the same order on every rank (a stream that shuffles needs a fixed seed of its own). A streaming loader with worker processes is refused: a dataset that splits itself across workers, as a Hugging Face streaming dataset does, would lose samples silently. It is not main's vlm_dataset_utils._ShardedIterable, which pads each rank to a count its loader knows: this wraps any loader's stream, which may have no length.
• distributed_forward_loop builds the forward_loop that mtq.quantize calls. It moves each batch to the device it is given and runs the full distributed forward, so the expert-parallel collectives fire inside it, with KV caching off as on the single-process path. Each batch is a mapping of the forward's keyword arguments; anything else is refused with a TypeError at the first batch.
  • A model whose distribution plan ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736) is expert-parallel passes the ep_mesh its experts are sharded over, and the loop re-plants the torch_a2a forward wrappers that quantizer insertion clobbers on transformers >= 5.16. A module whose forward is still the wrapper planted there is skipped, so passes do not stack wrappers; one whose forward was rebound since (another conversion, an export) is planted again. The planted wrappers stay on the model, so its later forwards (a preview's, an export's) work once the loop has run. mtq.auto_quantize calls its forward_step directly and would plant nothing, so it refuses a model whose recorded plan shards its experts, before converting anything. Any other model passes none. A mesh missing for an expert-parallel model, or given for another, raises.
  • It returns the allocator's reserved pool before the first forward and after the last. NCCL's buffers live outside the torch allocator, so a full reserved pool starves them. A process that never started CUDA skips this, and the garbage collection it costs.
• Every rank runs the same batch size. Equal step counts need equal batch sizes, and hf_ptq's default --batch_size 0 sizes calibration with get_max_batch_size on each rank: it measured every visible GPU, and on an out-of-memory error halved the batch and probed again, a different number of collective-bearing forwards per rank. For a model whose recorded plan spreads it over several ranks it is now collective: each rank measures the GPU it computes on (torch.cuda.current_device()) from one forward of the single-sample input, without probing the size it arrives at (a rank running out of memory partway through a probe's forward would leave its peers waiting in that forward's collectives), and every rank takes the smallest size any rank allows, reduced with all_reduce(..., min) ([OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736). A single process, and a model with no plan recorded (a Megatron model, say), are unchanged.
• Calibration that keeps its statistics on the rank is refused. Every rank calibrates on its own data, so an algorithm must share what it gathers across the ranks or the ranks calibrate their copies, or their shards, of a weight differently. max all-reduces its amax, GPTQ merges its Hessians across data-parallel replicas, and MSE and SmoothQuant work from max's synced amax. AWQ-Lite does not: it averages its activation scales over data parallelism but picks each layer's scale from its own rank's search, and a 2-rank probe showed the replicas' weights and pre_quant_scale diverging. awq_clip now averages each candidate clip's loss over data parallelism before the ranks pick one, in both branches: each rank accumulates its own losses, and after the forward loop every module's losses and token count are all-reduced once, in module order, so the collective count does not depend on how many forwards a rank ran a module (an MoE expert's is each rank's own). Its static-block branch averaged per forward before, and its per-tensor branch, which NVFP4 takes, not at all; this also makes Megatron data-parallel replicas agree on the per-tensor clip. local_hessian and imatrix share nothing either, and neither do max and MSE with distributed_sync off. mtq.quantize, mtq.calibrate and mtq.auto_quantize refuse these (awq_lite, awq_full, svdquant, which runs awq_lite, local_hessian, imatrix, and any config that turns distributed_sync off, with False or None, also as another algorithm's weight_scale_algorithm or scale_algorithm) on a model whose recorded plan spans several ranks, before converting anything (auto_quantize in its format check). A model with no plan recorded, such as a Megatron model, is unchanged. On an expert-parallel plan, any calibration is refused unless every module holding routed experts carries, as its expert-parallel group, the rank's row of the (edp, ep) mesh the expert-parallel load lays the ranks out in, ep innermost (in row order, which picks each rank's expert shard), and, as its data-parallel group, its column, the plan's edp ranks holding the same experts: each rank's per-expert quantizers hold its own experts, and max's and GPTQ's data-parallel sync would otherwise merge different experts' scales. The modules holding routed experts are those carrying an expert-parallel group and, once quantization has converted them, the quantized modules under a container named experts, as transformers' expert-parallel plans name it; unseeded experts beside a seeded router are refused once converted, before they calibrate. [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741's expert-parallel loader gives the routed experts those groups. A module with no routed experts and no expert-parallel group, such as a VLM's vision tower that hf_ptq quantizes on its own with every quantizer disabled, passes.
• hf_ptq's --use_fsdp2 path is not changed here. Its load (parallel_load_and_prepare_fsdp2) records no plan, so the batch sizing and the refusal above do not apply to it yet. Recording the plan there lands with [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it, and export an expert-parallel model through it #2752, the first PR whose export reads it on that path, together with stopping a planned run from retrying a batch after a rank-local out-of-memory error; [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754 then moves the load onto load_pretrained, which records the plan itself. nvfp4_awq_clip's data-parallel clip (below) needs no plan, so it applies to --use_fsdp2 runs now.

shard_dataloader refuses a loader that picks its own samples (a custom batch_sampler, a sampler other than sequential, random or distributed, or one of those over another source than the loader's dataset, also inside a plain BatchSampler or with auto-batching off) instead of rebuilding it, and keeps a plain BatchSampler's batch size and drop_last in the rank's shard; a loader with auto-batching off (batch_size=None) shards one item per step, as it loads. With one rank, shard_dataloader returns the loader itself, its sampler and shuffling included. distributed_forward_loop runs its forwards without gradients. get_max_batch_size reads the recorded plan, as the calibration guard does. distributed_forward_loop refuses an ep_mesh whose size is not the plan's ep_size.

Usage

from modelopt.torch.utils.distributed import distributed_forward_loop, rank, shard_dataloader, size

loader = shard_dataloader(calib_loader, rank(), size())
# ep_mesh: the mesh an expert-parallel model's experts are sharded over; None for any other model
forward_loop = distributed_forward_loop(model, loader, device=device, ep_mesh=ep_mesh)
model = mtq.quantize(model, quant_cfg, forward_loop=forward_loop)

Testing

• New tests/unit/torch/utils/test_distributed_calibration.py (42 cases), on CPU, in one process except where noted. It covers:

  • strided shards that cover the data exactly once, with equal batch counts on uneven splits, for map-style and streaming datasets; a streaming loader with workers refused; map-style and streaming loaders with auto-batching off sharded one item per step;
  • every loader setting carried over, a plain BatchSampler's batch size and drop_last included;
  • the wrapper re-planting, on the planned modules only, without stacking, and again after another conversion drops it;
  • the loop running this rank's shard through the top-level model, with use_cache off and the allocator cache freed around it in a process that started CUDA, and not in one that did not; a batch that is not a mapping refused;
  • the ep_mesh refused when missing for an expert-parallel plan, given for another, or over the whole world for an ep_size=2 plan;
  • rank-local calibration refused through quantize, calibrate and auto_quantize (before conversion), also when another algorithm runs it for its weight scales or max is told not to sync, with the model left as it was;
  • on two gloo ranks, max calibration on an expert-parallel plan run when the experts carry their groups, and refused with no expert-parallel group anywhere or with the experts' data-parallel group over every rank (before conversion), and with the group on the router only (once the experts are converted, before they calibrate);
  • a module without routed experts calibrating under an expert-parallel plan, with every quantizer disabled (hf_ptq's vision-tower call) and quantized.

• distributed_forward_loop runs on real GPU ranks in [OMNIML-6061] Load the dense part sharded, and size the load warmup to each rank #2742's expert-parallel calibration test, which checks each rank's per-expert scales against the whole model; this PR's own tests run it in one process.

• tests/unit/torch/utils/test_dataset_utils.py, on two real gloo ranks (cpu_dist_workers) with only the CUDA memory queries faked per rank: each rank measures only its own GPU, probes once, and both take the smaller batch size. A rank keeping its own size, measuring every GPU, or skipping the reduction fails it, the last after the group's 10 s timeout rather than by hanging; main's single-process retry test still passes. The two-rank test also runs a model with no plan recorded, which each rank measures on every visible GPU, as a single process does, without a refusal.

• Mutation-checked: deciding the reduction or the mesh other than by the plan fails a test (5 mutants), and so does dropping either refusal call, the nested-algorithm check, the distributed_sync check, hf_ptq's plan recording, the expert-parallel group refusal, or the ep_mesh size check. Refusing every expert-parallel calibration, or reading the plan strictly in get_max_batch_size, fails a test too. Dropping the check on converted experts or the data-parallel size check, keeping a replant flag once set, or planting on every pass each fail a test. Rebuilding a lone rank's loader, refusing a module without routed experts, or passing any module with no expert-parallel group each fail a test too.

• The shard tests are one table over dataset kind and world size, the two replant tests are one, and cases a remaining test covers are dropped. Each drop is checked by mutation.

• A subset sampler is refused, also under a plain BatchSampler and with auto-batching off, and so is a custom batch sampler; a lone rank keeps each such loader as it is. The sampler, batching, non-mapping-batch and auto_quantize cases each fail without their fix (mutation-checked).

• Locally on this revision, the stack rebased on main c6fa613: pre-commit over the whole PR against main, at every PR of the stack; and at the stack's tip the full tests/unit on transformers 5.18 and 5.5 and on torch 2.8, and the hf_ptq example tests, each with the same failures as main in that environment. CI runs the full unit suite.

• GPU tests were not re-run on this revision.

• Re-cut on 2026-10-27 (see the note at the top): pre-commit, mypy included, passes at this head over the whole change against main, and the series' PRs merged together give a tree byte-identical to the former stack's tip. This PR's own unit tests pass at this head on CPU (torch 2.14, transformers 5.18); CI runs the full unit suite.

• Review fixes, 2026-10-27: the a2a forward wrappers are re-planted only where quantizer insertion rebound the forward, so a wrapper still in place is never wrapped twice (found by review: the double wrap ran a module's input transforms twice, which no real run has shown to fail; on 4 GPUs at transformers 5.16.1 the expert-parallel load tests pass with either behaviour, and the CPU calibration-loop test is what tells them apart); the install probe checks each style it calls. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

• Review fixes, 2026-10-27, round 2: the batch size is measured on the GPU the rank computes on (torch.cuda.current_device()); a streaming loader with workers is refused; the docstring and the hf_ptq README say why the measurement does not probe. Each behaviour fix has a test that fails without it; pre-commit, mypy included, and this PR's own unit tests pass at this head.

• Rebased on 2026-10-27 onto main 7c3c868, after [OMNIML-6064] Load each rank's checkpoint shard straight into a meta-built model #2739 merged: pre-commit, mypy included, and this PR's own unit tests pass at this head, and the series' PRs merged together add no unit-test failure over main (tests/unit/torch utils, export, quantization and models, on transformers 5.18 and 5.5.4 and on torch 2.8). The hf_ptq README's FSDP2 AutoQuantize note follows main's KV-cache AutoQuantize support ([OMNIML-5987] [2/3] Support DP/FSDP2 KV AutoQuant and checkpoint resume #2664): weight and KV-cache AutoQuantize with --use_fsdp2 are experimental, and refused when a candidate format is one of those this PR refuses.

• Review fix, 2026-10-27 night: shard_dataloader refuses a RandomSampler that does not draw the whole dataset once (other num_samples, or with replacement), as it refuses other samplers.

• Review suggestion, 2026-10-27, last round: the comment over the sync guard names the algorithms' defaults it relies on.

• Rebased on 2026-10-28, late night, onto [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 at cf4597dcbeb, on main 1ca8bd1fc4d (after [OMNIML-6073] [1/3] Gather only FSDP2's own parameters in its weight-access window #3088 and [OMNIML-6073] [2/3] Quantize fused experts under grouped_mm; leave transposed layouts unquantized #2745 merged), without conflicts; this PR's own diff is unchanged. Its tests with [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736's and the hf_ptq argument tests pass on transformers 5.18 (220), and pre-commit passes against main.

• Rebased on 2026-10-28, late night, onto main bdc45368d10 after [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736 merged; [OMNIML-6052] Add the distribution plan, a run-scoped process group, and collectives that fail together #2736's squash has the tree of its last head, so this PR's tree is unchanged.

• Review fixes, 2026-10-29: the two-rank batch-size test makes its bounded group with bounded_collectives, whose barrier keeps a rank from destroying the group while its peer is still connecting (one CI run had failed so); the batch sizing's distributed helpers are imported at module top; the rank-local refusal reads each algorithm's options through its mode's schema, so a distributed_sync of "false" is refused; the experts' groups are checked by the ranks they span, an expert-parallel group of ep_size ranks and a data-parallel group of edp ranks meeting it only at the rank itself (a test answers four ranks' group queries, which the two-rank pool cannot build). Each behaviour fix fails its test without it (mutation-checked). Pre-commit passes against main; this PR's tests, test_distributed.py, test_quantize_cpu.py and the hf_ptq argument tests pass on CPU with transformers 5.18 and 5.5.4.

• Review fix, 2026-10-29, second round: awq_clip's per-tensor branch averages each candidate clip's loss over data parallelism, through the helper its static-block branch now uses too (same group, same op, every rank in the same order), and awq_clip is allowed on a distributed model; this also makes Megatron data-parallel replicas agree on the per-tensor clip. A two-rank CPU test of the static-block branch fails without the averaging (mutation-checked); a two-GPU NVFP4 test of the per-tensor branch is in tests/gpu (NVFP4 fake quantization needs CUDA; not run here). The hf_ptq README and the changelog say nvfp4_awq_clip works with --use_fsdp2 on several GPUs. Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU.

• GPU check, 2026-10-29, of the NVFP4 per-tensor branch on GB200 (torch 2.11.0a0, transformers 5.9.0), with the second round's per-forward averaging; without the data-parallel mean the even and odd ranks pick different per-tensor clips:

  Run	1 GPU	2 GPUs	4 GPUs
  test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank	skipped (needs 2+)	passed, 23.2 s	passed, 24.1 s
  The same test with the averaging reverted	–	failed: clips 95 vs 100	–

• GPU check, 2026-10-29, at cdb4e952468, with the fourth round's once-per-module averaging after the forward loop, on GB200 (torch 2.11.0a0, transformers 5.9.0); with that averaging disabled, the even and odd ranks pick different per-tensor clips again:

  Run at cdb4e952468	1 GPU	2 GPUs	4 GPUs
  test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank	skipped (needs 2+)	passed, 20.7 s	passed, 24.5 s
  The same test with sync_across_dp disabled (return as its first statement)	–	failed: clips 95 vs 100	–

• GPU check, 2026-10-29, at 01a1df717e3, with the fifth round's empty-input guard, on GB200 (torch 2.11.0a0, transformers 5.9.0). The test now also runs a second linear on no tokens on the odd ranks; without the guard, that linear's clip stays 0 on every replica:

  Run at 01a1df717e3	1 GPU	2 GPUs	4 GPUs
  test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank	skipped (needs 2+)	passed, 21.1 s	passed, 25.0 s
  The same test without the empty-input guard in _clip_search	–	failed on amax[1] > 0: amax = [100, 0]	–

• Review fixes, 2026-10-29, third round: a distributed_sync of None, which skips the sync as False does, is refused too, nested scale algorithms included; the experts' groups are checked against the (edp, ep) mesh, the expert-parallel group as the rank's row in row order and the data-parallel group as its column, so diagonal data-parallel groups and a reordered expert-parallel group are refused (the four-rank test adds both). Each fails its test without the change (mutation-checked). Pre-commit passes against main; this PR's tests pass on CPU.

• Review fixes, 2026-10-29, fourth round: awq_clip averages its clip search over the data-parallel ranks once per module after the forward loop, with the token count, in both branches, instead of once per forward, which hung when ranks ran an MoE expert a different number of times (the two-rank test now does; it fails without the change, on its bounded group's timeout) (the same round accepted the experts' groups of a mesh with either axis innermost; the eighth round reverts that). Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU, and the quantization unit tests pass with the change.

• Review fix, 2026-10-29, fifth round: awq_clip skips its clip search for an input with no tokens, which an MoE expert a rank routes none to can still get. Before, the per-tensor branch added NaN losses, which the post-loop sum spread to the replicas that had tokens (a clip of 0 for them), and the static-block branch failed to reshape the empty input. The two distributed tests now run an expert on no tokens on one rank, and check that its clips are positive and those a single process picks from the other ranks' tokens; each fails without the change (mutation-checked; the GPU one in the table above). Pre-commit passes against main; test_calib.py and this PR's tests pass on CPU.

• Review fix, 2026-10-29, sixth round: hf_ptq's --use_fsdp2 plan recording, its README notes, changelog entry and argument test leave this PR for [OMNIML-6083] Write a sharded model's HF checkpoint without gathering it, and export an expert-parallel model through it #2752 (see above); the new GPU test tests/gpu/torch/quantization/test_fsdp2.py::test_planned_batch_size_then_max_calibration (2+ GPUs) sizes an FSDP2 model's batch by its recorded plan with the real get_max_batch_size, rank 0's own estimate far above the others', checks every rank takes the smallest, then calibrates each rank's shard with max through distributed_forward_loop and checks the steps and scales agree. Its logic ran on two CPU ranks (FSDP2 on gloo, the CUDA memory queries faked): it passes, and without the batch size's reduction it fails at once on the sizes ([40, 4]); GPU results in the table below. Rebased onto main 2b900c8fa41; this PR's own diff is otherwise unchanged. Pre-commit passes against main.

• GPU check, 2026-10-29, at 10b3c1f0e35, of the new FSDP2 test on GB200 (torch 2.11.0a0, CUDA 13.1, NCCL 2.29.2, transformers 5.9.0); without the cross-rank minimum, rank 0 keeps its own estimate while the others take theirs:

  Run at 10b3c1f0e35	1 GPU	2 GPUs	4 GPUs
  test_fsdp2.py::test_planned_batch_size_then_max_calibration	skipped (needs 2+)	passed, 22.3 s	passed, 25.0 s
  The same test without the minimum's all_reduce in the planned branch	–	failed: sizes [36, 4]	–

• Review fix, 2026-10-29, seventh round: stacked on [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 (rebased from main 2b900c8fa41 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's ec9e09c3fc0; this PR's own diff is unchanged). New GPU test tests/gpu/torch/utils/test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model (2+ GPUs, transformers ≥ 5.16): a tiny Qwen3-MoE loaded by transformers' own expert-parallel load on the torch_a2a styles and experts forward, quantized with FP8 through distributed_forward_loop, each rank on its own batch; every expert's weight scales and every activation scale match the whole model calibrated on every rank's batch, with bounded collectives. Its logic ran on two CPU ranks with transformers 5.18: it passes, re-planting the forward wrappers quantizer insertion drops, and fails without the re-plant, or with the conversion forcing eager as before [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074. GPU results in the table below. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 and main.

• GPU check, 2026-10-29, at d90a479b6b5, of the new expert-parallel test on GB200 (torch 2.11.0a0, CUDA 13.1, NCCL 2.29.2, transformers 5.18.0): the expert-parallel calibration matches the whole model's, and needs both the forward re-plant and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074:

  Run at d90a479b6b5	1 GPU	2 GPUs	4 GPUs
  test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model	skipped	passed, 21.6 s	passed, 26.7 s
  The same test with the re-plant disabled (_reinstall_a2a_forwards returns 0)	–	failed: "torch_a2a experts forward requires the EP process_group"	–
  The same test without [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 (normalize_experts_impl_on_the_fly forcing eager, main's behaviour)	–	failed: "Qwen3MoeExperts.forward() takes 4 positional arguments but 5 were given"	–

• Restacked on 2026-10-29 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 at f68837967d3, now [2/2] on [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 [1/2] after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's split; no conflicts, and this PR's own diff is unchanged. The new expert-parallel test's logic still passes on two CPU ranks on this base. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's head.

• GPU check, 2026-10-29, at 54f0c1fd8e9 (on [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 f68837967d3, on [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 33736db61fa), on GB200 (torch 2.11.0a0, transformers 5.18.0): the expert-parallel calibration still matches the whole model with [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122's per-rank expert calibration underneath:

  Run at 54f0c1fd8e9	2 GPUs	4 GPUs
  test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model	passed, 23.5 s	passed, 27.0 s

• Review fixes, 2026-10-29, eighth round (Claude's review of 54f0c1fd8e9): the experts' groups are checked against the ep-innermost (edp, ep) mesh only, the layout the expert-parallel load builds: accepting either axis innermost let transposed groups pass when edp == ep_size (a new four-rank case, groups-transposed, fails with that check). awq_clip reduces each module's clip search in one all-reduce, the token count riding in the losses' buffer (the ninth round reduces the losses in place instead), and every module's reduction runs before any weight gather. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074; test_calib.py and this PR's tests pass on CPU, and dropping the token count's sync fails the two-rank test.

• On GPUs, with the rest of the expert-parallel series on top (transformers 5.13), before this round's change to awq_clip:

  • Qwen3-30B-A3B at EP 4 and EP 2 on 4 GPUs, deterministic, batch size 1: every exported scale is bit-identical between the two. In the EP 2 run (edp 2), the same row/column check, reading the groups from the field that stack's loader seeds, accepted the loader's real (edp, ep) groups.

    Exported scale	Differing, EP 4 vs EP 2
    weight_scale	0 of 18,624
    input_scale	0 of 18,624
    k_scale	0 of 48
    v_scale	0 of 48

  • Qwen3-0.6B with nvfp4_awq_clip and --use_fsdp2 on 2 GPUs: awq_clip ran and the sharded checkpoint was exported, end to end. That run had the per-forward averaging this round replaces.

• Review fixes, 2026-10-29, ninth round (Claude's review of 2dfeb9c23f0): awq_clip's candidate losses are views of one buffer, which the sync all-reduces and averages in place, with the token count in its own int64 all-reduce: no copy of the losses, where the packed buffer held up to three, and every module's reductions still run before any weight gather. The refusal names distributed_sync only for max and MSE, the calibrations that have it. Pre-commit passes against [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074; test_calib.py and this PR's tests pass on CPU, and the two-rank test fails without either all-reduce, or with the candidates' losses copied out of the buffer.

• Review follow-up, 2026-10-29, tenth round (Claude's review of b517c7581b6): a docstring fix (07affa0fe76). The new expert-parallel test's job, run on two CPU ranks with transformers 5.18, shows that its dense linears are QuantLinear with the seeded data-parallel group and no tensor-parallel group (Qwen3-MoE's expert plan shards no dense module). Seeding them without that group fails the test: o_proj's input amax is off by 0.293 (relative).

• Review follow-up, 2026-10-29, eleventh round (Claude's review of 07affa0fe76): a docstring note (e546bcca0b1) that every rank still reads and decodes every sample of a sharded stream. The experts' check that quantize() runs before conversion reads the groups a loader seeds. [OMNIML-6060] [3/3] Refuse an expert-parallel load that would calibrate the experts wrong #2741's load is the first to seed them; until it lands, the unit tests and the expert-parallel GPU test, which seed them by hand, exercise that check.

• Rebased on 2026-10-30 onto main 99c3834f62c through [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 1e74c31b49a and [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 3f3544435d8 (from [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's f5d36d4f794), as fd036d44f28; no conflicts, and this PR's own diff is unchanged. Pre-commit passes against main.

• GPU check, 2026-10-30, at fd036d44f28, on GB200 (torch 2.11.0a0, transformers 5.9.0): awq_clip's in-place data-parallel sync (the ninth round's change) still gives every rank the same clip:

  Run at fd036d44f28	2 GPUs	4 GPUs
  test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank	passed, 18 s	passed, 24 s

• Restacked on 2026-10-30 after [OMNIML-6073] [1/2] Calibrate and fold expert-parallel fused experts on each rank's own experts #3122 merged: onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 39697ad985b, which is on main 55898ee4e3b (from 3f3544435d8), as 3f56acfbb76; no conflicts, and this PR's own diff is unchanged. Pre-commit passes against main. This PR's unit tests (test_distributed_calibration.py, test_calib.py, test_dataset_utils.py, test_distributed.py) pass on CPU with transformers 5.18 and 5.5.4, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

• Rebased on 2026-10-30 onto [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 a6758d77eb3, on main 059588b89a5, which has [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128 (from 3f56acfbb76), as 186efae02e5. [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's AWQ guard (_refuse_awq_on_a_sharded_fsdp2_model) and this PR's (_refuse_unsupported_distributed_calibration) met at the same three call sites in model_quant.py: calibrate(), quantize() and auto_quantize's candidate check. Both stay, [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's first. The changelog and the README keep both sides, and 186efae02e5 makes this PR's nvfp4_awq_clip README note its own paragraph, saying why it is not refused. This PR's own diff is otherwise unchanged. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py (with [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's guard test) and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

• Rebased on 2026-10-30 onto main d464d30b120 after [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074 merged ([OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's squash; from 186efae02e5), as 9a87d153c89: [OMNIML-6073] [2/2] Calibrate fused experts on grouped_mm and torch_a2a instead of forcing eager #3074's own commits drop out; no conflicts, and this PR's own diff is unchanged, with [OMNIML-6081] Fix the FSDP2 export of W4A8 (SequentialQuantizer) models, and refuse compressing them #3128's guard still first at the three call sites. This PR now targets main, so the full CI runs, GPU included. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

• Review fixes, 2026-10-30, twelfth round (cjluo-nv's review of 9a87d153c89): mtq.auto_quantize refuses a model whose recorded plan shards its experts, before converting it (fbbe2275cf8); a two-rank test with seeded experts and an INT8 candidate checks the module types and state are unchanged, and fails without the refusal. shard_dataloader compares a sampler's source with the loader's dataset (e1976572627, typed in 9ac1486cdab); a SequentialSampler, RandomSampler or DistributedSampler over a shorter source is refused, and the three new cases fail with the old check. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

• Review fix, 2026-10-31, thirteenth round (Claude's review of 9ac1486cdab): the forward re-plant now plants only over a bound method, the forward quantizer insertion leaves, as its docstring says (aba16882660); a planned module whose forward the load never wrapped is left alone, which a new test checks and the old guard fails. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py and this PR's unit tests pass, and the expert-parallel GPU test's job still re-plants the two experts' wrappers and passes on two CPU ranks with transformers 5.18.

• Rebased on 2026-10-31 onto main e352caca3c3 (from aba16882660, adding Run fused-MoE expert GPTQ updates on all of this process's GPUs #3002, Collect imatrix statistics with one fused float32 reduction #3044 and [OMNIML-6273] Let the Claude review bot submit its review again #3156), as dea9b793553; no conflicts (Collect imatrix statistics with one fused float32 reduction #3044 and Run fused-MoE expert GPTQ updates on all of this process's GPUs #3002 touch model_calib.py, which merges cleanly), and this PR's own diff is unchanged. Pre-commit passes against main. On CPU with transformers 5.18 and 5.5.4, test_quantize_cpu.py, the ggml GPTQ and imatrix unit tests and this PR's unit tests pass, and the expert-parallel GPU test's job passes on two CPU ranks with transformers 5.18.

• Review follow-up, 2026-10-31, fourteenth round (Claude's review of dea9b793553): a unit test checks that every calibration _RANK_LOCAL_CALIBRATION refuses names a registered calibrate mode, so a renamed or misspelled one fails instead of silently losing its refusal (d10be569b04; it fails with one name misspelled); 7791f6d0f3b drops a redundant test import. The sampler check keeps comparing lengths (see its thread). CI's GPU job passed at dea9b793553; the three multi-GPU tests skip on its single-GPU runner, and their run at 7791f6d0f3b on 2 and 4 GPUs is below. Pre-commit passes against main; the changed test files pass on CPU with transformers 5.18 and 5.5.4.

• GPU check, 2026-10-31, at 7791f6d0f3b, on 2 and 4 GPUs (torch 2.15, transformers 5.18.0): all three multi-GPU tests pass, none skipped. Call times; the first test also pays the shared worker pool's startup, so times do not compare across tests:

  Run at 7791f6d0f3b	2 GPUs	4 GPUs
  test_calib_cuda.py::test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank	passed, 12.0 s	passed, 13.7 s
  test_fsdp2.py::test_planned_batch_size_then_max_calibration	passed, 1.3 s	passed, 3.1 s
  test_moe_a2a.py::test_expert_parallel_calibration_matches_the_whole_model	passed, 3.9 s	passed, 10.1 s

• Trimmed on 2026-10-31 (95148c3ff8b), after a review counted the PR over the ~500-line budget: four docstrings and a comment that restated the code are shorter, with no code change. 492 added source lines (git diff --shortstat e352caca3c3...95148c3ff8b -- . ':!tests' ':!docs'). Pre-commit passes against main, and the unit tests pass on CPU with transformers 5.18 and 5.5.4.

• Rebased on 2026-11-01 onto main cf38ec30ca4 (from 95148c3ff8b), as ff372fbc30c: no conflicts, range-diff 63 of 63 commits unchanged; pre-commit with all hooks, mypy included, passes against main; on CPU, 272 unit tests pass with transformers 5.18 (the expert-parallel job on two CPU ranks included) and 271 with 5.5.4.

• Rebased on 2026-11-01 onto main 51645db66dc after [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 merged (from ff372fbc30c), as 6a06a5421f7: no conflicts, range-diff 60 of 63 commits unchanged and 3 differing only in a context line [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 changed; pre-commit with all hooks, mypy included, passes against main; on CPU, 269 unit tests pass with transformers 5.18 (the expert-parallel job on two CPU ranks included) and 268 with 5.5.4, three fewer than before because [OMNIML-6054] [1/2] Move fsdp2_wrap out of distributed.py into fsdp2_utils #3140 moved the fsdp2_wrap tests to their own file.

• Rebased on 2026-11-02 onto main 2d4e8d019e7 after [OMNIML-6054] [2/2] Route load_pretrained by the distribution plan #2779 and [OMNIML-6076] [2/3] Prepare a model's config for expert parallelism from its spec #3135 merged (from 6a06a5421f7), as 5829a854612: no conflicts, range-diff 63 of 63 commits unchanged, 492 added source lines; pre-commit with all hooks, mypy included, passes against main; on CPU, 269 unit tests pass with transformers 5.18 (the expert-parallel job on two CPU ranks included) and 268 with 5.5.4.

• Review fixes, 2026-11-02 (meenchen's reviews of 5829a854612): each calibrate mode descriptor declares _distributed_calibration (True for none, max, mse, smoothquant, awq_clip, gptq, nvfp4_act_headroom and lsq; False for awq_lite, awq_full, svdquant, local_hessian and imatrix) and checks its own options and nested algorithms; a mode that declares neither is refused under a multi-rank plan, and _RANK_LOCAL_CALIBRATION is gone (c6585dce94d, 49911504cf4). The 13 built-in modes refuse what they did before; a test checks every mode modelopt registers declares (it fails with one declaration removed), another that an undeclared mode is refused (it fails when undeclared modes pass). e8d80fed068 keeps the get_max_batch_size and shard_dataloader docstrings to their contract, with the reasons moved to comments. Pre-commit with all hooks passes against main; on CPU, 273 unit tests pass with transformers 5.18 (the expert-parallel job on two CPU ranks included) and 272 with 5.5.4. Claude's review of 49911504cf4 approved.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ Additive: get_max_batch_size and the calibration refusals change only for a model whose recorded plan spans several ranks. awq_clip's per-tensor (NVFP4) clip changes under data parallelism, where its replicas now agree (changelog, Bug Fixes).
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅ A Bug Fixes entry for AWQ-Clip's data-parallel clip; the feature's entry lands with [OMNIML-6085] [3/3] Run expert-parallel PTQ from hf_ptq.py #2754.
• Did you get Claude approval on this PR?: ❌ Not yet.

Additional Information

Part 8 of 21 of expert-parallel PTQ for Hugging Face MoE models; merges after #2739, which has merged.

Summary by CodeRabbit

• New Features
  • Distributed calibration now shards map-style and streaming datasets across ranks, keeping batch counts aligned and preserving data loader settings.
  • Calibration supports expert- and tensor-parallel models with improved device handling, cache management, and compatibility with parallel model wrappers.
  • Distributed batch-size detection probes each rank’s local GPU and selects the smallest supported batch size across ranks.
• Bug Fixes
  • Calibration and quantization now reject selected rank-local calibration methods on distributed models before modifying model state.